using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp;
using Microsoft.CodeAnalysis.CSharp.Syntax;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace NetCorePal.Extensions.CodeAnalysis.SourceGenerators;

[Generator]
public class CommandMetadataGenerator : IIncrementalGenerator
{
    public void Initialize(IncrementalGeneratorInitializationContext context)
    {
        var typeDeclarations = context.SyntaxProvider
            .CreateSyntaxProvider(
                predicate: (node, _) => node is ClassDeclarationSyntax || node is RecordDeclarationSyntax,
                transform: (ctx, _) => ctx.Node)
            .Where(n => n is ClassDeclarationSyntax || n is RecordDeclarationSyntax);

        var compilationAndTypes = context.CompilationProvider.Combine(typeDeclarations.Collect());

        context.RegisterSourceOutput(compilationAndTypes, (spc, source) =>
        {
            var (compilation, typeNodes) = ((Compilation, System.Collections.Immutable.ImmutableArray<SyntaxNode>))source;
            var commandTypes = new HashSet<string>();

            foreach (var typeDecl in typeNodes)
            {
                var semanticModel = compilation.GetSemanticModel(typeDecl.SyntaxTree);
                var symbol = semanticModel.GetDeclaredSymbol(typeDecl) as INamedTypeSymbol;
                if (symbol == null) continue;
                if (GeneratorExtensions.IsCommand(symbol))
                {
                    commandTypes.Add(symbol.ToDisplayString());
                }
            }

            if (commandTypes.Count > 0)
            {
                var sb = new StringBuilder();
                sb.AppendLine("// <auto-generated/>\nusing System;\nusing NetCorePal.Extensions.CodeAnalysis.Attributes;");
                foreach (var cmd in commandTypes)
                {
                    sb.AppendLine($"[assembly: CommandMetadataAttribute(\"{cmd}\")]\n");
                }
                spc.AddSource("CommandMetadata.g.cs", sb.ToString());
            }
        });
    }
}

